package com.tbryant.test.redis.usagescenarios;

import com.google.common.collect.Maps;
import com.tbryant.test.redis.RedisConstant;
import redis.clients.jedis.Jedis;

import java.util.Map;

/**
 * 分布式锁
 * 支持可重入
 * TODO：不能保证del key的线程 = set key的线程
 * TODO：竞争失败没有重试机制
 */
public class DistributedLock {
    private static final Jedis JEDIS = new Jedis(RedisConstant.REDIS_IP, RedisConstant.REDIS_PORT);
    private static ThreadLocal<Map<String, Integer>> threadLocal = new ThreadLocal<>();

    static boolean lock(String key) {
        Map<String, Integer> threadLocalMap = getThreadLocalMap();
        Integer refCount = threadLocalMap.get(key);
        // 不为空表示已获得锁
        if (null != refCount) {
            // 计数+1，不操作redis
            threadLocalMap.put(key, refCount + 1);
            return true;
        }
        boolean result = doLock(key);
        if (!result) {
            return false;
        }
        // 首次加锁，计数1
        threadLocalMap.put(key, 1);
        return true;
    }

    static boolean unlock(String key) {
        Map<String, Integer> threadLocalMap = getThreadLocalMap();
        Integer refCount = threadLocalMap.get(key);
        // 为空表示未获得锁
        if (null == refCount) {
            return false;
        }
        refCount -= 1;
        if (refCount > 0) {
            // 仍持有锁，更新计数，不操作redis
            threadLocalMap.put(key, refCount);
        } else {
            // 释放锁，只移除当前key的计数
            threadLocalMap.remove(key);
            doUnLock(key);
        }
        return true;
    }

    private static Map<String, Integer> getThreadLocalMap() {
        Map<String, Integer> threadLocalMap = threadLocal.get();
        if (null != threadLocalMap) {
            return threadLocalMap;
        }
        threadLocal.set(Maps.newHashMap());
        return threadLocal.get();
    }

    /**
     * 操作redis竞争锁
     */
    private static boolean doLock(String key) {
        /**
         * setnx + expire
         * key不存在，返回OK
         * key已存在，返回null
         */
        return JEDIS.set(key, RedisConstant.REDIS_LOCK, RedisConstant.REDIS_NX, RedisConstant.REDIS_EX, RedisConstant.DISTRIBUTED_LOCK_EX_TIME) != null;
    }

    /**
     * 操作redis释放锁
     */
    private static void doUnLock(String key) {
        JEDIS.del(key);
    }
}